
import csv
from utils.dataloader import get_data



#-------------------------------------#
#   设置数据集的路径
#   train_txt_path   训练图片路径和标签
#   valid_txt_path   验证图片路径和标签
#-------------------------------------#
train_txt_path   = "cls_train.txt"
valid_txt_path   = "cls_valid.txt"


if __name__ == "__main__":

    #----------------------#
    #   读取数据集对应的txt
    #----------------------#
    with open(train_txt_path, "r") as f:
        train_lines = f.readlines()
    with open(valid_txt_path, "r") as f:
        val_lines = f.readlines()

    # 获取训练数据和验证数据的数量
    num_train = len(train_lines)
    num_val = len(val_lines)

    print(f"训练样本数量:{num_train}, 验证样本数量:{num_val}, 计算图像特征量中...")
    train_data = get_data(train_lines)
    valid_data = get_data(val_lines)

    print("将图像特征量和标签写入csv文件")
    with open("cls_train.csv", 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        title = ["image mean", "fire mea", "fire area", "label"]
        writer.writerow(title)
        for row in train_data:
            writer.writerow(row)

    with open("cls_valid.csv", 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        title = ["image mean", "fire mea", "fire area", "label"]
        writer.writerow(title)
        for row in valid_data:
            writer.writerow(row)


